#ifndef AVL_TREE_H
#define AVL_TREE_H

#include "../bit_operations/bit_operator.h"

template <typename K, typename V>
class AvlTree {
	enum Direction { STAY = -1, LEFT = 0, RIGHT = 1 };
	struct node {
		node(K k, V value) : value(value), lean(0) { key[0] = k; children[LEFT] = children[RIGHT] = 0x0; }

		node* children[2];
		V value;
		K key[1];
		char lean;

		inline char Lean() { return lean; }

		inline bool Lean(char offset) {
			bool r = (lean == offset);
			lean += offset;
			return r;
		}

		inline void SetLean(char l) { lean = l; }

		Direction RelativeDirection(const K &k) {
			comparison_t t;
			if (!BitOperator<K>::Compare(key, &k, t)) return STAY;
			return (t < 0) ? RIGHT : LEFT;
		}

		inline Direction HeavierDirection() { return (lean == 1) ? RIGHT : LEFT; }

		inline void OverwriteWith(const node *other) {
			children[LEFT] = other->children[LEFT];
			children[RIGHT] = other->children[RIGHT];
			lean = other->lean;
		}

		void RecursiveDelete() { 
			children[LEFT]->RecursiveDelete(); delete children[LEFT];
			children[RIGHT]->RecursiveDelete(); delete children[RIGHT];
		}
	};


	struct Step {
		node* element;
		Direction d;
		Step() {}
		Step(node* e, Direction d) : element(e), d(d) {}
		inline bool InsertUpdate() { return element->Lean(mapDtoL[d]); }
		inline bool RemoveUpdate() { return element->Lean(-mapDtoL[d]); }
		inline node** ChildReference() { return &(element->children[d]); }
	};

	inline static void Rotate(node** root, Direction d, Direction ad) {
		node *n = *root, *t = n->children[ad];
		n->children[ad] = t->children[d];
		t->children[d] = n;
		*root = t;
	}

	void InsertionRotate(node** root, Step &s) {
		Direction d = s.d, ad = (Direction)!d;
		node *element = s.element, *child = element->children[d];
		int_fast8_t l = mapDtoL[d], al = -l, cl = child->Lean();

		if (cl == l) {
			element->SetLean(0);
			child->SetLean(0);
		} else {
			node *childchild = child->children[ad];
			int_fast8_t ccl = childchild->Lean();

			if (!ccl) {
				element->SetLean(0);
				child->SetLean(0);
			} else if (ccl == al) {
				element->SetLean(0);
				child->SetLean(l);
			} else {
				element->SetLean(al);
				child->SetLean(0);
			}

			childchild->SetLean(0);

			Rotate(&(element->children[d]), d, ad);
		}

		Rotate(root, ad, d);
	}

	bool RemovalRotate(node** root, Step &s) {
		Direction ad = s.d, d = (Direction)!ad;
		node *element = s.element, *child = element->children[d];
		int_fast8_t l = mapDtoL[d], al = -l, cl = child->Lean();

		bool balanced = false;
		if (cl == 0) {
			element->SetLean(l);
			child->SetLean(al);
			balanced = true;
		} else if (cl == l) {
			element->SetLean(0);
			child->SetLean(0);
		} else {
			node *childchild = child->children[ad];
			int_fast8_t ccl = childchild->Lean();

			if (!ccl) {
				element->SetLean(0);
				child->SetLean(0);
			} else if (ccl == al) {
				element->SetLean(0);
				child->SetLean(l);
			} else {
				element->SetLean(al);
				child->SetLean(0);
			}

			childchild->SetLean(0);

			Rotate(&(element->children[d]), d, ad);
		}

		Rotate(root, ad, d);

		return balanced;

	}

	inline void PathReset() { current = 0; }
	inline void PathPush(node* n, Direction d) { path[current++] = Step(n, d); }
	inline void PathPop() { --current; }

	Step path[sizeof(K)* 8];
	node* root = 0x0;
	uint_least32_t current = 0;
	const static int_fast8_t mapDtoL[2];

public:
	~AvlTree() { root->RecursiveDelete(); delete root; }
	void Insert(const K &key, const V &value) {
		if (!root) root = new node(key, value);
		else {
			PathReset();

			//Find insertion location and build the traversal path
			Direction d;
			node *context = root;
			for (; context && (d = context->RelativeDirection(key)) != STAY; context = context->children[d]) PathPush(context, d);
			//If the location found has an element, then the key already exists
			if (context) return;
			//Put a new node into the location
			unsigned int i = current - 1;
			Step s = path[i];

			s.element->children[d] = new node(key, value);

			//Go backwards through the path
			for (; i > 0; s = path[--i]) {
				//If the update calls for a rebalance, do the rotate and break
				if (s.InsertUpdate()) { InsertionRotate(path[i-1].ChildReference(), s); break; } 
				//If the update balances this node, then everything above it won't need to update
				else if (!s.element->Lean()) break;
			}

			//Check the top of the tree for an update last, if it needs one
			if (i == 0 && s.InsertUpdate()) InsertionRotate(&root, s);
		}
	}
	void Remove(const K &key) {
		if (!root) return;

		PathReset();

		Direction d;
		node *context = root;
		for (; context && (d = context->RelativeDirection(key)) != STAY; context = context->children[d]) PathPush(context, d);

		if (!context) return;
		node **reference = (current == 0 ? &root : path[current-1].ChildReference());
		
		unsigned int i = current;
		Direction heavydirection = context->HeavierDirection();
		PathPush(context, heavydirection);
		
		context = context->children[heavydirection];
		if (context) {	
			for (d = (Direction)!heavydirection; context; context = context->children[d]) PathPush(context, d);

			context = path[current - 1].element;
			Step s = path[current - 2];
			s.element->children[s.d] = context->children[heavydirection];
			context->OverwriteWith(*reference);
			path[i].element = context;
		}

		delete *reference; *reference = context;

		if (--current == 0) return;

		//Rebalance up the path
		i = current - 1;
		Step s = path[i];
		//Go backwards through the path
		for (; i > 0; s = path[--i]) {
			//If the update calls for a rebalance, do the rotate and break if the rotation didn't change sub-tree height
			if (s.RemoveUpdate()) {
				if (RemovalRotate(path[i - 1].ChildReference(), s)) break;
				//If the update balances this node, then everything above it won't need to update
			} else if (s.element->Lean()) break;
		}

		//Check the top of the tree for an update last, if it needs one
		if (i == 0 && s.RemoveUpdate()) RemovalRotate(&root, s);
	}
	V* Find(const K &key) {
		node* context = root;
		for (Direction d; context && (d = context->RelativeDirection(key)) != STAY; context = context->children[d]);

		if (context) return &(context->value);
		else return 0x0;
	}
	V* operator [](const K &key) { return Find(key); }
};

template <typename K, typename V>
const int_fast8_t AvlTree<K, V>::mapDtoL[2] = { -1, 1 };

#endif